from pupper.hardware_interface import HardwareInterface
from pupper.inverse_kinematics import four_legs_inverse_kinematics
from pupper.controller import Controller
from pupper.config import Command, State, BehaviorState, Configuration
import time
import numpy as np
from transforms3d.euler import quat2euler
from ahrs.filters import Mahony
from pi_driver import LepiDriver


def main():
    """
    主函数
    """

    # 初始化各种对象
    config = Configuration()
    controller = Controller(
        config,
    )
    state = State()
    state.behavior_state = BehaviorState.REST
    command = Command()

    hardware_interface = HardwareInterface()

    # 初始姿态
    state.foot_locations = (
        config.default_stance
        + np.array([0, 0, command.height])[:, np.newaxis]
    )
    last_loop = time.time()

    print("Summary of gait parameters:")
    print("overlap time: ", config.overlap_time)
    print("swing time: ", config.swing_time)
    print("z clearance: ", config.z_clearance)
    print("x shift: ", config.x_shift)

    lepi = LepiDriver()
    orientation = Mahony(frequency=100)
    Q_ = np.array([1., 0., 0., 0.])

    while True:
        now = time.time()
        # 每过一个时钟周期才进行一次计算
        if now - last_loop < config.dt:
            continue
        last_loop = time.time()

        acc = np.array(lepi.readAccData(), dtype=np.float)
        gyro = np.array(lepi.readGyroData(), dtype=np.float)
        # magn = np.array(lepi.readMagnData(), dtype=np.float)
        # 进行姿态解算, 地磁数据未修正前误差较大所以不予计算
        Q = orientation.updateIMU(
            Q_, gyr=gyro*250/32768/180*np.pi, acc=acc*9.8/16384)
        Q_ = Q
        euler = np.array(quat2euler(Q))
        pitch, roll, yaw = euler/np.pi*180

        # 旋转补偿
        if abs(pitch) > 5:
            if abs(command.pitch) <= config.max_pitch:
                command.pitch -= 0.01 * pitch/abs(pitch)
            else:
                command.pitch = config.max_pitch * pitch/abs(pitch)

        if abs(roll) > 5:
            if abs(command.roll) <= config.max_roll:
                command.roll += 0.01 * roll/abs(roll)
            else:
                command.roll = config.max_roll * roll/abs(roll)
        # 执行一次坐标计算
        controller.run(state, command)

        # 打印计算出来的坐标
        # print(state.rotated_foot_locations)
        print(pitch, roll)
        # 坐标到关节角度
        joint_angles = four_legs_inverse_kinematics(
            state.rotated_foot_locations, config)

        # 关节角度到舵机指令
        hardware_interface.set_actuator_postions(joint_angles)


if __name__ == '__main__':
    main()
